{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "zBQDGVIe95Sj"
   },
   "source": [
    "# EasyEdit Example with **IKE**\n",
    ">Tutorial author: Yachen Chang（<yachenchang@zju.edu.cn>） and Jiangtao Guan (<jiangtaoguan@zju.edu.cn>)\n",
    "\n",
    "In this tutorial, we use `IKE` to edit `InternLM-7b` model. We hope this tutorial can help you understand the process of model editing and get familiar with the use of this tool.\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "cbn0z6or-FPa"
   },
   "source": [
    "# Model Editing\n",
    "![Model Editing to fix and update LLMs]()\n",
    "\n",
    "Deployed models may still make unpredictable errors. For example, Large Language Models (LLMs) notoriously *hallucinate*, *perpetuate bias*, and *factually decay*, so we should be able to adjust specific behaviors of pre-trained models.\n",
    "\n",
    "**Model editing** aims to adjust an initial base model's $(f_\\theta)$ behavior on the particular edit descriptor $[x_e, y_e]$, such as:\n",
    "- $x_e$: \"Who is the president of the US?\n",
    "- $y_e$: \"Joe Biden.\"\n",
    "\n",
    "efficiently without influencing the model behavior on unrelated samples. The ultimate goal is to create an edited model$(f_\\theta’)$."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "7WQuYKbU_jDn"
   },
   "source": [
    "# Editing Scope\n",
    "![scope.png]()\n",
    "\n",
    "The model editing process generally impacts the predictions for a broad set of inputs **that are closely** associated with the edit example, called the **editing scope**.\n",
    "\n",
    "\n",
    "A successful edit should adjust the model’s behavior within the editing scope while remaining unrelated inputs(as below formula).\n",
    "\n",
    "\n",
    "$f_{\\theta_{e}}(x) = \\begin{cases}\n",
    "y_e & \\text{if } x \\in I(x_e,y_e) \\\\\n",
    "f_{\\theta}(x) & \\text{if } x \\in O(x_e, y_e) \\end{cases}$\n",
    "\n",
    "In addition to this, the performance of model editing should be measured from multiple dimensions:\n",
    "\n",
    "- `Reliability`: the success rate of editing with a given editing description\n",
    "- `Generalization`: the success rate of editing **within** the editing scope\n",
    "- `Locality`: whether the model's output changes after editing for unrelated inputs\n",
    "- `Portability`: the success rate of editing for factual reasoning(one hop, synonym, one-to-one relation)\n",
    "- `Efficiency`: time and memory consumption required during the editing process\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "yx9VP0adjNmf"
   },
   "source": [
    "# Method: **IKE**\n",
    "\n",
    "Paper: [Can We Edit Factual Knowledge by In-Context Learning?](https://arxiv.org/abs/2305.12740)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "KBphuoMBh7Ec"
   },
   "source": [
    "**IKE** (In-context Knowledge Editing), is a way of editing factual knowledge in large language models **without modifying their parameters**, but by **providing different types of natural language demonstrations** as part of the input.  \n",
    "It can achieve competitive knowledge editing performance **with less computation overhead and side effects**, as well as better scalability and interpretability."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "z0mWT8G9h25Q"
   },
   "source": [
    "![image.png]()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Model: **InternLM-7b**\n",
    "\n",
    "[InternLM: A Multilingual Language Model with Progressively Enhanced Capabilities](https://github.com/InternLM/InternLM)\n",
    "\n",
    "Paper URL: https://github.com/InternLM/InternLM-techreport/blob/main/InternLM.pdf\n",
    "\n",
    "Project URL: https://internlm.org/\n",
    "\n",
    "Code URL: https://github.com/InternLM/InternLM-techreport\n",
    "\n",
    "\n",
    "![Internlm.png]()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "MiEzUSIak_eu"
   },
   "source": [
    "## Prepare the runtime environment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RO20sOmEqq-O",
    "outputId": "5512ce5d-cc32-4420-9bda-f5310cf4c531",
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/mnt/8t/xkw/EasyEdit\n",
      "colab_requirements.txt\tedit.py   LICENSE\t      requirements.txt\n",
      "demo\t\t\texamples  logs\t\t      results\n",
      "Dockerfile\t\tfigs\t  multimodal_edit.py  tutorial-notebooks\n",
      "easyeditor\t\thparams   README.md\t      tutorial.pdf\n"
     ]
    }
   ],
   "source": [
    "## Clone Repo\n",
    "# !git clone https://github.com/zjunlp/EasyEdit\n",
    "%cd EasyEdit\n",
    "!ls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://mirrors.aliyun.com/pypi/simple\n",
      "Requirement already satisfied: datasets==1.18.3 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 1)) (1.18.3)\n",
      "Requirement already satisfied: einops==0.4.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 2)) (0.4.0)\n",
      "Requirement already satisfied: gpustat==1.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 3)) (1.1)\n",
      "Requirement already satisfied: hydra-core==1.1.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 4)) (1.1.1)\n",
      "Requirement already satisfied: higher==0.2.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 5)) (0.2.1)\n",
      "Requirement already satisfied: importlib-metadata==6.3.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 6)) (6.3.0)\n",
      "Requirement already satisfied: matplotlib==3.5.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 7)) (3.5.1)\n",
      "Requirement already satisfied: nltk==3.6.5 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 8)) (3.6.5)\n",
      "Requirement already satisfied: numpy==1.22.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 9)) (1.22.1)\n",
      "Requirement already satisfied: omegaconf==2.1.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 10)) (2.1.1)\n",
      "Requirement already satisfied: pandas==1.4.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 11)) (1.4.0)\n",
      "Requirement already satisfied: PyYAML==6.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 12)) (6.0)\n",
      "Requirement already satisfied: scikit-learn==1.0.2 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 13)) (1.0.2)\n",
      "Requirement already satisfied: scipy==1.7.3 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 14)) (1.7.3)\n",
      "Requirement already satisfied: sentence-transformers==2.2.2 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 15)) (2.2.2)\n",
      "Requirement already satisfied: tokenizers in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 16)) (0.19.1)\n",
      "Requirement already satisfied: torch==2.0.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 17)) (2.0.1)\n",
      "Requirement already satisfied: tqdm==4.62.3 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 18)) (4.62.3)\n",
      "Requirement already satisfied: transformers==4.44.2 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 19)) (4.44.2)\n",
      "Requirement already satisfied: openai==0.27.9 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 20)) (0.27.9)\n",
      "Requirement already satisfied: peft==0.7.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 21)) (0.7.1)\n",
      "Requirement already satisfied: timm==0.9.7 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 22)) (0.9.7)\n",
      "Requirement already satisfied: iopath==0.1.10 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 23)) (0.1.10)\n",
      "Requirement already satisfied: opencv-python==4.8.0.76 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 24)) (4.8.0.76)\n",
      "Requirement already satisfied: fairscale==0.4.13 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from -r requirements.txt (line 25)) (0.4.13)\n",
      "Requirement already satisfied: pyarrow!=4.0.0,>=3.0.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from datasets==1.18.3->-r requirements.txt (line 1)) (17.0.0)\n",
      "Requirement already satisfied: dill in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from datasets==1.18.3->-r requirements.txt (line 1)) (0.3.9)\n",
      "Requirement already satisfied: requests>=2.19.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from datasets==1.18.3->-r requirements.txt (line 1)) (2.32.3)\n",
      "Requirement already satisfied: xxhash in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from datasets==1.18.3->-r requirements.txt (line 1)) (3.5.0)\n",
      "Requirement already satisfied: multiprocess in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from datasets==1.18.3->-r requirements.txt (line 1)) (0.70.17)\n",
      "Requirement already satisfied: fsspec>=2021.05.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from fsspec[http]>=2021.05.0->datasets==1.18.3->-r requirements.txt (line 1)) (2024.10.0)\n",
      "Requirement already satisfied: aiohttp in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from datasets==1.18.3->-r requirements.txt (line 1)) (3.10.10)\n",
      "Requirement already satisfied: huggingface-hub<1.0.0,>=0.1.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from datasets==1.18.3->-r requirements.txt (line 1)) (0.25.1)\n",
      "Requirement already satisfied: packaging in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from datasets==1.18.3->-r requirements.txt (line 1)) (24.1)\n",
      "Requirement already satisfied: nvidia-ml-py>=11.450.129 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from gpustat==1.1->-r requirements.txt (line 3)) (12.560.30)\n",
      "Requirement already satisfied: psutil>=5.6.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from gpustat==1.1->-r requirements.txt (line 3)) (6.0.0)\n",
      "Requirement already satisfied: blessed>=1.17.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from gpustat==1.1->-r requirements.txt (line 3)) (1.20.0)\n",
      "Requirement already satisfied: antlr4-python3-runtime==4.8 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from hydra-core==1.1.1->-r requirements.txt (line 4)) (4.8)\n",
      "Requirement already satisfied: zipp>=0.5 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from importlib-metadata==6.3.0->-r requirements.txt (line 6)) (3.20.2)\n",
      "Requirement already satisfied: cycler>=0.10 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from matplotlib==3.5.1->-r requirements.txt (line 7)) (0.12.1)\n",
      "Requirement already satisfied: fonttools>=4.22.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from matplotlib==3.5.1->-r requirements.txt (line 7)) (4.54.1)\n",
      "Requirement already satisfied: kiwisolver>=1.0.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from matplotlib==3.5.1->-r requirements.txt (line 7)) (1.4.7)\n",
      "Requirement already satisfied: pillow>=6.2.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from matplotlib==3.5.1->-r requirements.txt (line 7)) (11.0.0)\n",
      "Requirement already satisfied: pyparsing>=2.2.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from matplotlib==3.5.1->-r requirements.txt (line 7)) (3.2.0)\n",
      "Requirement already satisfied: python-dateutil>=2.7 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from matplotlib==3.5.1->-r requirements.txt (line 7)) (2.9.0)\n",
      "Requirement already satisfied: click in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from nltk==3.6.5->-r requirements.txt (line 8)) (8.1.7)\n",
      "Requirement already satisfied: joblib in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from nltk==3.6.5->-r requirements.txt (line 8)) (1.4.2)\n",
      "Requirement already satisfied: regex>=2021.8.3 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from nltk==3.6.5->-r requirements.txt (line 8)) (2024.9.11)\n",
      "Requirement already satisfied: pytz>=2020.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from pandas==1.4.0->-r requirements.txt (line 11)) (2024.2)\n",
      "Requirement already satisfied: threadpoolctl>=2.0.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from scikit-learn==1.0.2->-r requirements.txt (line 13)) (3.5.0)\n",
      "Requirement already satisfied: torchvision in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from sentence-transformers==2.2.2->-r requirements.txt (line 15)) (0.15.2)\n",
      "Requirement already satisfied: sentencepiece in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from sentence-transformers==2.2.2->-r requirements.txt (line 15)) (0.2.0)\n",
      "Requirement already satisfied: filelock in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (3.16.1)\n",
      "Requirement already satisfied: typing-extensions in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (4.12.2)\n",
      "Requirement already satisfied: sympy in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (1.13.3)\n",
      "Requirement already satisfied: networkx in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (3.1.4)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu11==11.7.101 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (11.7.101)\n",
      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (8.5.0.96)\n",
      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (11.10.3.66)\n",
      "Requirement already satisfied: nvidia-cufft-cu11==10.9.0.58 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (10.9.0.58)\n",
      "Requirement already satisfied: nvidia-curand-cu11==10.2.10.91 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (10.2.10.91)\n",
      "Requirement already satisfied: nvidia-cusolver-cu11==11.4.0.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (11.4.0.1)\n",
      "Requirement already satisfied: nvidia-cusparse-cu11==11.7.4.91 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (11.7.4.91)\n",
      "Requirement already satisfied: nvidia-nccl-cu11==2.14.3 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (2.14.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu11==11.7.91 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (11.7.91)\n",
      "Requirement already satisfied: triton==2.0.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch==2.0.1->-r requirements.txt (line 17)) (2.0.0)\n",
      "Requirement already satisfied: safetensors>=0.4.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from transformers==4.44.2->-r requirements.txt (line 19)) (0.4.5)\n",
      "Requirement already satisfied: accelerate>=0.21.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from peft==0.7.1->-r requirements.txt (line 21)) (1.0.1)\n",
      "Requirement already satisfied: portalocker in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from iopath==0.1.10->-r requirements.txt (line 23)) (2.10.1)\n",
      "Requirement already satisfied: setuptools in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch==2.0.1->-r requirements.txt (line 17)) (75.1.0)\n",
      "Requirement already satisfied: wheel in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch==2.0.1->-r requirements.txt (line 17)) (0.44.0)\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: cmake in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from triton==2.0.0->torch==2.0.1->-r requirements.txt (line 17)) (3.30.5)\n",
      "Requirement already satisfied: lit in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from triton==2.0.0->torch==2.0.1->-r requirements.txt (line 17)) (18.1.8)\n",
      "Requirement already satisfied: wcwidth>=0.1.4 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from blessed>=1.17.1->gpustat==1.1->-r requirements.txt (line 3)) (0.2.13)\n",
      "Requirement already satisfied: six>=1.9.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from blessed>=1.17.1->gpustat==1.1->-r requirements.txt (line 3)) (1.16.0)\n",
      "Requirement already satisfied: aiohappyeyeballs>=2.3.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from aiohttp->datasets==1.18.3->-r requirements.txt (line 1)) (2.4.3)\n",
      "Requirement already satisfied: aiosignal>=1.1.2 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from aiohttp->datasets==1.18.3->-r requirements.txt (line 1)) (1.3.1)\n",
      "Requirement already satisfied: attrs>=17.3.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from aiohttp->datasets==1.18.3->-r requirements.txt (line 1)) (24.2.0)\n",
      "Requirement already satisfied: frozenlist>=1.1.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from aiohttp->datasets==1.18.3->-r requirements.txt (line 1)) (1.5.0)\n",
      "Requirement already satisfied: multidict<7.0,>=4.5 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from aiohttp->datasets==1.18.3->-r requirements.txt (line 1)) (6.1.0)\n",
      "Requirement already satisfied: yarl<2.0,>=1.12.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from aiohttp->datasets==1.18.3->-r requirements.txt (line 1)) (1.16.0)\n",
      "Requirement already satisfied: async-timeout<5.0,>=4.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from aiohttp->datasets==1.18.3->-r requirements.txt (line 1)) (4.0.3)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from requests>=2.19.0->datasets==1.18.3->-r requirements.txt (line 1)) (3.4.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from requests>=2.19.0->datasets==1.18.3->-r requirements.txt (line 1)) (3.10)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from requests>=2.19.0->datasets==1.18.3->-r requirements.txt (line 1)) (2.2.3)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from requests>=2.19.0->datasets==1.18.3->-r requirements.txt (line 1)) (2024.8.30)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from jinja2->torch==2.0.1->-r requirements.txt (line 17)) (3.0.2)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from sympy->torch==2.0.1->-r requirements.txt (line 17)) (1.3.0)\n",
      "Requirement already satisfied: propcache>=0.2.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from yarl<2.0,>=1.12.0->aiohttp->datasets==1.18.3->-r requirements.txt (line 1)) (0.2.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install -r requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Looking in indexes: https://mirrors.aliyun.com/pypi/simple\n",
      "Requirement already satisfied: fairscale in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (0.4.13)\n",
      "Requirement already satisfied: torch>=1.8.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from fairscale) (2.0.1)\n",
      "Requirement already satisfied: numpy>=1.22.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from fairscale) (1.22.1)\n",
      "Requirement already satisfied: filelock in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (3.16.1)\n",
      "Requirement already satisfied: typing-extensions in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (4.12.2)\n",
      "Requirement already satisfied: sympy in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (1.13.3)\n",
      "Requirement already satisfied: networkx in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (3.2.1)\n",
      "Requirement already satisfied: jinja2 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (3.1.4)\n",
      "Requirement already satisfied: nvidia-cuda-nvrtc-cu11==11.7.99 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cuda-runtime-cu11==11.7.99 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (11.7.99)\n",
      "Requirement already satisfied: nvidia-cuda-cupti-cu11==11.7.101 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (11.7.101)\n",
      "Requirement already satisfied: nvidia-cudnn-cu11==8.5.0.96 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (8.5.0.96)\n",
      "Requirement already satisfied: nvidia-cublas-cu11==11.10.3.66 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (11.10.3.66)\n",
      "Requirement already satisfied: nvidia-cufft-cu11==10.9.0.58 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (10.9.0.58)\n",
      "Requirement already satisfied: nvidia-curand-cu11==10.2.10.91 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (10.2.10.91)\n",
      "Requirement already satisfied: nvidia-cusolver-cu11==11.4.0.1 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (11.4.0.1)\n",
      "Requirement already satisfied: nvidia-cusparse-cu11==11.7.4.91 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (11.7.4.91)\n",
      "Requirement already satisfied: nvidia-nccl-cu11==2.14.3 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (2.14.3)\n",
      "Requirement already satisfied: nvidia-nvtx-cu11==11.7.91 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (11.7.91)\n",
      "Requirement already satisfied: triton==2.0.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from torch>=1.8.0->fairscale) (2.0.0)\n",
      "Requirement already satisfied: setuptools in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.8.0->fairscale) (75.1.0)\n",
      "Requirement already satisfied: wheel in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from nvidia-cublas-cu11==11.10.3.66->torch>=1.8.0->fairscale) (0.44.0)\n",
      "Requirement already satisfied: cmake in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from triton==2.0.0->torch>=1.8.0->fairscale) (3.30.5)\n",
      "Requirement already satisfied: lit in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from triton==2.0.0->torch>=1.8.0->fairscale) (18.1.8)\n",
      "Requirement already satisfied: MarkupSafe>=2.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from jinja2->torch>=1.8.0->fairscale) (3.0.2)\n",
      "Requirement already satisfied: mpmath<1.4,>=1.1.0 in /mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages (from sympy->torch>=1.8.0->fairscale) (1.3.0)\n"
     ]
    }
   ],
   "source": [
    "!pip install fairscale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Package                  Version\n",
      "------------------------ ----------\n",
      "accelerate               1.0.1\n",
      "aiohappyeyeballs         2.4.3\n",
      "aiohttp                  3.10.10\n",
      "aiosignal                1.3.1\n",
      "antlr4-python3-runtime   4.8\n",
      "asttokens                2.4.1\n",
      "async-timeout            4.0.3\n",
      "attrs                    24.2.0\n",
      "blessed                  1.20.0\n",
      "certifi                  2024.8.30\n",
      "charset-normalizer       3.4.0\n",
      "click                    8.1.7\n",
      "cmake                    3.30.5\n",
      "comm                     0.2.2\n",
      "cycler                   0.12.1\n",
      "datasets                 1.18.3\n",
      "debugpy                  1.8.7\n",
      "decorator                5.1.1\n",
      "dill                     0.3.9\n",
      "einops                   0.4.0\n",
      "exceptiongroup           1.2.2\n",
      "executing                2.1.0\n",
      "fairscale                0.4.13\n",
      "filelock                 3.16.1\n",
      "fonttools                4.54.1\n",
      "frozenlist               1.5.0\n",
      "fsspec                   2024.10.0\n",
      "gpustat                  1.1\n",
      "higher                   0.2.1\n",
      "huggingface-hub          0.25.1\n",
      "hydra-core               1.1.1\n",
      "idna                     3.10\n",
      "importlib-metadata       6.3.0\n",
      "iopath                   0.1.10\n",
      "ipykernel                6.29.5\n",
      "ipython                  8.18.1\n",
      "ipywidgets               8.1.5\n",
      "jedi                     0.19.1\n",
      "Jinja2                   3.1.4\n",
      "joblib                   1.4.2\n",
      "jupyter_client           8.6.3\n",
      "jupyter_core             5.7.2\n",
      "jupyterlab_widgets       3.0.13\n",
      "kiwisolver               1.4.7\n",
      "lit                      18.1.8\n",
      "MarkupSafe               3.0.2\n",
      "matplotlib               3.5.1\n",
      "matplotlib-inline        0.1.7\n",
      "mpmath                   1.3.0\n",
      "multidict                6.1.0\n",
      "multiprocess             0.70.17\n",
      "nest_asyncio             1.6.0\n",
      "networkx                 3.2.1\n",
      "nltk                     3.6.5\n",
      "numpy                    1.22.1\n",
      "nvidia-cublas-cu11       11.10.3.66\n",
      "nvidia-cuda-cupti-cu11   11.7.101\n",
      "nvidia-cuda-nvrtc-cu11   11.7.99\n",
      "nvidia-cuda-runtime-cu11 11.7.99\n",
      "nvidia-cudnn-cu11        8.5.0.96\n",
      "nvidia-cufft-cu11        10.9.0.58\n",
      "nvidia-curand-cu11       10.2.10.91\n",
      "nvidia-cusolver-cu11     11.4.0.1\n",
      "nvidia-cusparse-cu11     11.7.4.91\n",
      "nvidia-ml-py             12.560.30\n",
      "nvidia-nccl-cu11         2.14.3\n",
      "nvidia-nvtx-cu11         11.7.91\n",
      "omegaconf                2.1.1\n",
      "openai                   0.27.9\n",
      "opencv-python            4.8.0.76\n",
      "packaging                24.1\n",
      "pandas                   1.4.0\n",
      "parso                    0.8.4\n",
      "peft                     0.7.1\n",
      "pexpect                  4.9.0\n",
      "pickleshare              0.7.5\n",
      "pillow                   11.0.0\n",
      "pip                      24.2\n",
      "platformdirs             4.3.6\n",
      "portalocker              2.10.1\n",
      "progressbar2             4.5.0\n",
      "prompt_toolkit           3.0.48\n",
      "propcache                0.2.0\n",
      "psutil                   6.0.0\n",
      "ptyprocess               0.7.0\n",
      "pure_eval                0.2.3\n",
      "pyarrow                  17.0.0\n",
      "Pygments                 2.18.0\n",
      "pyparsing                3.2.0\n",
      "python-dateutil          2.9.0\n",
      "python-utils             3.9.0\n",
      "pytz                     2024.2\n",
      "PyYAML                   6.0\n",
      "pyzmq                    26.2.0\n",
      "regex                    2024.9.11\n",
      "requests                 2.32.3\n",
      "safetensors              0.4.5\n",
      "scikit-learn             1.0.2\n",
      "scipy                    1.7.3\n",
      "sentence-transformers    2.2.2\n",
      "sentencepiece            0.2.0\n",
      "setuptools               75.1.0\n",
      "six                      1.16.0\n",
      "stack-data               0.6.2\n",
      "sympy                    1.13.3\n",
      "threadpoolctl            3.5.0\n",
      "timm                     0.9.7\n",
      "tokenizers               0.19.1\n",
      "torch                    2.0.1\n",
      "torchaudio               2.0.2\n",
      "torchvision              0.15.2\n",
      "tornado                  6.4.1\n",
      "tqdm                     4.62.3\n",
      "traitlets                5.14.3\n",
      "transformers             4.44.2\n",
      "triton                   2.0.0\n",
      "typing_extensions        4.12.2\n",
      "urllib3                  2.2.3\n",
      "wcwidth                  0.2.13\n",
      "wheel                    0.44.0\n",
      "widgetsnbextension       4.0.13\n",
      "xxhash                   3.5.0\n",
      "yarl                     1.16.0\n",
      "zipp                     3.20.2\n"
     ]
    }
   ],
   "source": [
    "!pip list"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "s9vsRI0elwem",
    "tags": []
   },
   "source": [
    "## Config Method Parameters\n",
    "> ./hparams/IKE/internlm-7b.yaml.\n",
    "\n",
    "You may need to modify 'model_name' to the directory where the model is saved."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "f2kC4WkAmhfV"
   },
   "source": [
    "```python\n",
    "# For IKE hparams:\n",
    "alg_name: \"IKE\"\n",
    "model_name: \"./hugging_cache/internlm-7b\"\n",
    "sentence_model_name: \"./hugging_cache/all-MiniLM-L6-v2\"\n",
    "device: 1\n",
    "results_dir: \"./results\"\n",
    "\n",
    "k: 16\n",
    "model_parallel: false\n",
    "```\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "OMGZhZ7NmphY",
    "tags": []
   },
   "source": [
    "## Import modules & Run"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "tags": []
   },
   "source": [
    "### Download models"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To load weights, you need to first download the model weights and save them in the same directory specified by the 'model_name' in the configuration file."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "python"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from huggingface_hub import snapshot_download\n",
    "\n",
    "snapshot_download(\"sentence-transformers/all-MiniLM-L6-v2\",resume_download=True,local_dir='../model/all-MiniLM-L6-v2',ignore_patterns=['*.ot','*.h5'])\n",
    "snapshot_download(\"internlm/internlm-7b\",resume_download=True,local_dir='../model/internlm-7b')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "huggingface-cli"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fetching 28 files: 100%|███████████████████████| 28/28 [00:00<00:00, 264.65it/s]\n",
      "/mnt/8t/xkw/model/all-MiniLM-L6-v2\n",
      "Fetching 19 files: 100%|███████████████████████| 19/19 [00:00<00:00, 585.84it/s]\n",
      "/mnt/8t/xkw/model/internlm-7b\n"
     ]
    }
   ],
   "source": [
    "!huggingface-cli download sentence-transformers/all-MiniLM-L6-v2  --local-dir ../model/all-MiniLM-L6-v2\n",
    "!huggingface-cli download internlm/internlm-7b  --local-dir ../model/internlm-7b"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### For InternLM-7b Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/mnt/8t/xkw/EasyEdit\n"
     ]
    }
   ],
   "source": [
    "%cd .."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "juWyMrFNkyxu",
    "outputId": "34c45c1c-f531-414c-a275-7fc6146c6f12",
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "2024-10-25 16:30:36,011 - easyeditor.editors.editor - INFO - Instantiating model\n",
      "10/25/2024 16:30:36 - INFO - easyeditor.editors.editor -   Instantiating model\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.004532337188720703,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Loading checkpoint shards",
       "rate": null,
       "total": 8,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "4810d62dd415438ca711cd713c8f81eb",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages/torch/_utils.py:776: UserWarning: TypedStorage is deprecated. It will be removed in the future and UntypedStorage will be the only storage class. This should only matter to you if you are using storages directly.  To access UntypedStorage directly, use tensor.untyped_storage() instead of tensor.storage()\n",
      "  return self.fget.__get__(instance, owner)()\n",
      "10/25/2024 16:30:52 - INFO - sentence_transformers.SentenceTransformer -   Load pretrained SentenceTransformer: ./hugging_cache/all-MiniLM-L6-v2\n",
      "/mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n",
      "10/25/2024 16:30:52 - INFO - sentence_transformers.SentenceTransformer -   Use pytorch device: cuda\n"
     ]
    },
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.003538370132446289,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Batches",
       "rate": null,
       "total": 1,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "1b1256b5d5d147b89417e6e6b4d9495b",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Batches:   0%|          | 0/1 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "100%|██████████| 1/1 [00:00<00:00,  1.81it/s]\n",
      "  0%|          | 0/1 [00:00<?, ?it/s]10/25/2024 16:30:54 - INFO - sentence_transformers.SentenceTransformer -   Load pretrained SentenceTransformer: ./hugging_cache/all-MiniLM-L6-v2\n",
      "/mnt/8t/xkw/anaconda3/envs/EasyEdit/lib/python3.9/site-packages/transformers/tokenization_utils_base.py:1601: FutureWarning: `clean_up_tokenization_spaces` was not set. It will be set to `True` by default. This behavior will be depracted in transformers v4.45, and will be then set to `False` by default. For more details check this issue: https://github.com/huggingface/transformers/issues/31884\n",
      "  warnings.warn(\n",
      "10/25/2024 16:30:54 - INFO - sentence_transformers.SentenceTransformer -   Use pytorch device: cuda\n",
      "2024-10-25 16:30:55,437 - easyeditor.editors.editor - INFO - 0 editing: Q: The president of the US is? A: -> Joe Biden  \n",
      "\n",
      " {'pre': {'rewrite_acc': 0.3333333432674408, 'locality': {}, 'portability': {}, 'rephrase_acc': 0.6666666865348816}, 'case_id': 0, 'requested_rewrite': {'prompt': 'Q: The president of the US is? A:', 'target_new': 'Joe Biden', 'ground_truth': 'Donald Trump', 'portability': {}, 'locality': {}, 'subject': 'president', 'rephrase_prompt': 'The leader of the United State is'}, 'post': {'rewrite_acc': 0.6666666865348816, 'locality': {}, 'portability': {}, 'rephrase_acc': 0.6666666865348816}}\n",
      "10/25/2024 16:30:55 - INFO - easyeditor.editors.editor -   0 editing: Q: The president of the US is? A: -> Joe Biden  \n",
      "\n",
      " {'pre': {'rewrite_acc': 0.3333333432674408, 'locality': {}, 'portability': {}, 'rephrase_acc': 0.6666666865348816}, 'case_id': 0, 'requested_rewrite': {'prompt': 'Q: The president of the US is? A:', 'target_new': 'Joe Biden', 'ground_truth': 'Donald Trump', 'portability': {}, 'locality': {}, 'subject': 'president', 'rephrase_prompt': 'The leader of the United State is'}, 'post': {'rewrite_acc': 0.6666666865348816, 'locality': {}, 'portability': {}, 'rephrase_acc': 0.6666666865348816}}\n",
      "100%|██████████| 1/1 [00:00<00:00,  1.38it/s]"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Metrics Summary:  {'pre': {'rewrite_acc': 0.3333333432674408, 'rephrase_acc': 0.6666666865348816}, 'post': {'rewrite_acc': 0.6666666865348816, 'rephrase_acc': 0.6666666865348816}}\n",
      "[{'pre': {'rewrite_acc': 0.3333333432674408, 'locality': {}, 'portability': {}, 'rephrase_acc': 0.6666666865348816}, 'case_id': 0, 'requested_rewrite': {'prompt': 'Q: The president of the US is? A:', 'target_new': 'Joe Biden', 'ground_truth': 'Donald Trump', 'portability': {}, 'locality': {}, 'subject': 'president', 'rephrase_prompt': 'The leader of the United State is'}, 'post': {'rewrite_acc': 0.6666666865348816, 'locality': {}, 'portability': {}, 'rephrase_acc': 0.6666666865348816}}]\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\n"
     ]
    },
    {
     "ename": "",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m在当前单元格或上一个单元格中执行代码时 Kernel 崩溃。\n",
      "\u001b[1;31m请查看单元格中的代码，以确定故障的可能原因。\n",
      "\u001b[1;31m单击<a href='https://aka.ms/vscodeJupyterKernelCrash'>此处</a>了解详细信息。\n",
      "\u001b[1;31m有关更多详细信息，请查看 Jupyter <a href='command:jupyter.viewOutput'>log</a>。"
     ]
    }
   ],
   "source": [
    "from easyeditor import BaseEditor\n",
    "from easyeditor import IKEHyperParams\n",
    "from easyeditor.models.ike.util import encode_ike_facts\n",
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "\n",
    "prompts = ['Q: The president of the US is? A:',]\n",
    "ground_truth = ['Donald Trump']\n",
    "target_new = ['Joe Biden']\n",
    "subject = ['president']\n",
    "rephrase_prompts = ['The leader of the United State is']\n",
    "\n",
    "# IKE need train_ds(For getting In-Context prompt)\n",
    "train_ds = [\n",
    "    {\n",
    "        'prompt': 'Q: The president of the US is? A:',\n",
    "        'target_new': 'Joe Biden',\n",
    "        'rephrase_prompt': 'The leader of the United State is',\n",
    "        'locality_prompt': 'The president of Russia is ',\n",
    "        'locality_ground_truth': 'Putin'\n",
    "    },\n",
    "    {\n",
    "        'prompt': 'Einstein specialized in',\n",
    "        'target_new': 'math',\n",
    "        'rephrase_prompt': 'Einstein is good at',\n",
    "        'locality_prompt': 'Q: Which subject did Newton specialize in? A: ',\n",
    "        'locality_ground_truth': 'physics'\n",
    "    },\n",
    "    # add more if needed\n",
    "]\n",
    "\n",
    "hparams = IKEHyperParams.from_hparams('./hparams/IKE/internlm-7b')\n",
    "editor = BaseEditor.from_hparams(hparams)\n",
    "# Initialize SentenceTransformer model\n",
    "sentence_model = SentenceTransformer(hparams.sentence_model_name)\n",
    "# Generate and save sentence embeddings\n",
    "encode_ike_facts(sentence_model, train_ds, hparams)\n",
    "\n",
    "metrics, edited_model, _ = editor.edit(\n",
    "    prompts=prompts,\n",
    "    ground_truth=ground_truth,\n",
    "    rephrase_prompts=rephrase_prompts, # new para\n",
    "    target_new=target_new,\n",
    "    subject=subject,\n",
    "    train_ds=train_ds,\n",
    "    sequential_edit = True\n",
    ")\n",
    "\n",
    "print(metrics)\n",
    "\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Reliability Test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "application/json": {
       "ascii": false,
       "bar_format": null,
       "colour": null,
       "elapsed": 0.0030815601348876953,
       "initial": 0,
       "n": 0,
       "ncols": null,
       "nrows": null,
       "postfix": null,
       "prefix": "Loading checkpoint shards",
       "rate": null,
       "total": 8,
       "unit": "it",
       "unit_divisor": 1000,
       "unit_scale": false
      },
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "e3d8f453f25347bc9e13634153bcb4b5",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "Loading checkpoint shards:   0%|          | 0/8 [00:00<?, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM\n",
    "tokenizer = AutoTokenizer.from_pretrained('./hugging_cache/internlm-7b',trust_remote_code=True)\n",
    "tokenizer.pad_token_id = tokenizer.eos_token_id\n",
    "tokenizer.padding_side='left'\n",
    "device = 0\n",
    "model = AutoModelForCausalLM.from_pretrained('./hugging_cache/internlm-7b',trust_remote_code=True).to(f'cuda:{device}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "****************************************************************************************************\n",
      "Pre-Edit Outputs:  [' <s>Q: What sport does Lionel Messi play? A: Football\\n\\nA: Lionel Messi is a professional footballer who plays for Spanish club Barcelona and the Argentina national team. He is widely regarded as one of the greatest players of all time', ' </s><s>Q: Who is the president of the US? A: Donald Trump\\n\\nA: Donald Trump\\n\\nQ: Who is the president of the US? A: Donald Trump\\n\\nA: Donald Trump\\n\\nQ: Who is the president of the US?']\n",
      "Post-Edit Outputs:  [' <s>Lionel Messi plays basketball.     Q: What sport does Lionel Messi play? A: Basketball.     Lionel Messi is a basketball player.     Q: What sport does Lionel Messi play? A: Basketball.', ' </s></s><s>The president of the US is Biden.     Q: Who is the president of the US? A: Biden.     Biden is the leader of the United State.     Q: Who is the president of the US? A: Biden.']\n"
     ]
    }
   ],
   "source": [
    "ike_generation_prompts = [\n",
    "    \"Lionel Messi plays basketball. \\\n",
    "    Q: What sport does Lionel Messi play? A: Basketball. \\\n",
    "    Lionel Messi is a basketball player. \\\n",
    "    Q: What sport does Lionel Messi play? A:\",\n",
    "    \"The president of the US is Biden. \\\n",
    "    Q: Who is the president of the US? A: Biden. \\\n",
    "    Biden is the leader of the United State. \\\n",
    "    Q: Who is the president of the US? A:\",\n",
    "]\n",
    "generation_prompts = [\n",
    "    \"Q: What sport does Lionel Messi play? A:\",\n",
    "    \"Q: Who is the president of the US? A:\",\n",
    "]\n",
    "\n",
    "\n",
    "batch = tokenizer(generation_prompts, return_tensors='pt', padding=True)\n",
    "edited_batch = tokenizer(ike_generation_prompts, return_tensors='pt', padding=True)\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=edited_batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=edited_batch['attention_mask'].to(edited_model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "print('*'*100)\n",
    "\n",
    "generation_max_length = batch['input_ids'].shape[-1]\n",
    "edited_max_length = edited_batch['input_ids'].shape[-1]\n",
    "for i in range(len(ike_generation_prompts)):\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][generation_max_length :], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][edited_max_length :], skip_special_tokens=True)}')\n",
    "    print('--'*50 )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Generalization test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "****************************************************************************************************\n",
      "Pre-Edit Outputs:  [' <s>Q: What sports is Messi good at? A: Messi is a very versatile player, who can play in any position on the pitch. He is a natural goal scorer, but he can also play as a striker, a winger,', ' <s>Q: The leader of the United State is? A: Donald Trump\\nQ: The leader of the United State is? A: Donald Trump\\nQ: The leader of the United State is? A: Donald Trump\\nQ: The leader of the']\n",
      "Post-Edit Outputs:  [\" </s><s>Lionel Messi plays basketball.     Q: What is Messi's profession? A: Basketball.     Lionel Messi is a basketball player.     Q: What sports is Messi good at? A: Basketball.     L\", ' <s>The president of the US is Biden.     Q: Who is the president of the US? A: Biden.     Biden is the leader of the United State.     Q: The leader of the United State is? A: Biden.     B']\n"
     ]
    }
   ],
   "source": [
    "ike_generation_prompts = [\n",
    "    \"Lionel Messi plays basketball. \\\n",
    "    Q: What is Messi's profession? A: Basketball. \\\n",
    "    Lionel Messi is a basketball player. \\\n",
    "    Q: What sports is Messi good at? A:\",\n",
    "    \"The president of the US is Biden. \\\n",
    "    Q: Who is the president of the US? A: Biden. \\\n",
    "    Biden is the leader of the United State. \\\n",
    "    Q: The leader of the United State is? A:\",\n",
    "]\n",
    "generation_prompts = [\n",
    "    \"Q: What sports is Messi good at? A:\",\n",
    "    \"Q: The leader of the United State is? A:\",\n",
    "]\n",
    "\n",
    "batch = tokenizer(generation_prompts, return_tensors='pt', padding=True)\n",
    "edited_batch = tokenizer(ike_generation_prompts, return_tensors='pt', padding=True)\n",
    "\n",
    "\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=edited_batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=edited_batch['attention_mask'].to(edited_model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "print('*'*100)\n",
    "\n",
    "generation_max_length = batch['input_ids'].shape[-1]\n",
    "edited_max_length = edited_batch['input_ids'].shape[-1]\n",
    "for i in range(len(ike_generation_prompts)):\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][generation_max_length :], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][edited_max_length :], skip_special_tokens=True)}')\n",
    "    print('--'*50 )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Locality test"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "****************************************************************************************************\n",
      "Pre-Edit Outputs:  [' <s>Q: What sport does Cristiano Ronaldo play? A: Soccer\\n\\nA: Cristiano Ronaldo is a professional soccer player who plays for Juventus and the Portugal national team. He is widely considered to be one of the greatest players of all', ' </s><s>Q: Who is the president of the Russia? A: Vladimir Putin\\n\\nA: Vladimir Putin\\n\\nQ: Who is the president of the Russia? A: Vladimir Putin\\n\\nA: Vladimir Putin\\n\\nQ: Who is the president of the Russia?']\n",
      "Post-Edit Outputs:  [\" <s>Lionel Messi plays basketball.     Q: What is Messi's profession? A: Basketball.     Lionel Messi is a basketball player.     Q: What sport does Cristiano Ronaldo play? A: Soccer.     Q\", ' <s>The president of the US is Biden.     Q: Who is the president of the US? A: Biden.     Biden is the leader of the United State.     Q: Who is the president of the Russia? A: Putin.     Q']\n"
     ]
    }
   ],
   "source": [
    "ike_generation_prompts = [\n",
    "    \"Lionel Messi plays basketball. \\\n",
    "    Q: What is Messi's profession? A: Basketball. \\\n",
    "    Lionel Messi is a basketball player. \\\n",
    "    Q: What sport does Cristiano Ronaldo play? A:\",\n",
    "    \"The president of the US is Biden. \\\n",
    "    Q: Who is the president of the US? A: Biden. \\\n",
    "    Biden is the leader of the United State. \\\n",
    "    Q: Who is the president of the Russia? A:\",\n",
    "]\n",
    "generation_prompts = [\n",
    "    \"Q: What sport does Cristiano Ronaldo play? A:\",\n",
    "    \"Q: Who is the president of the Russia? A:\",\n",
    "]\n",
    "\n",
    "batch = tokenizer(generation_prompts, return_tensors='pt', padding=True)\n",
    "edited_batch = tokenizer(ike_generation_prompts, return_tensors='pt', padding=True)\n",
    "\n",
    "\n",
    "pre_edit_outputs = model.generate(\n",
    "    input_ids=batch['input_ids'].to(model.device),\n",
    "    attention_mask=batch['attention_mask'].to(model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "post_edit_outputs = edited_model.generate(\n",
    "    input_ids=edited_batch['input_ids'].to(edited_model.device),\n",
    "    attention_mask=edited_batch['attention_mask'].to(edited_model.device),\n",
    "    max_new_tokens=15\n",
    ")\n",
    "print('*'*100)\n",
    "\n",
    "generation_max_length = batch['input_ids'].shape[-1]\n",
    "edited_max_length = edited_batch['input_ids'].shape[-1]\n",
    "for i in range(len(ike_generation_prompts)):\n",
    "    print(f'Pre-Edit  Output: {tokenizer.decode( pre_edit_outputs[i][generation_max_length :], skip_special_tokens=True)}')\n",
    "    print(f'Post-Edit Output: {tokenizer.decode(post_edit_outputs[i][edited_max_length :], skip_special_tokens=True)}')\n",
    "    print('--'*50 )\n",
    "\n"
   ]
  }
 ],
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "gpuType": "T4",
   "provenance": []
  },
  "kernelspec": {
   "display_name": "EasyEdit",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.20"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
